from cvfo_mask import  CvFo
from tqdm import tqdm
import paddle
import numpy  as np
import pandas as pd
block_size =8
voc = pd.read_pickle("voc.pandas_pickle")
data = pd.read_pickle("dataset.pandas_pickle")
model = CvFo(len(voc), 512, 6,block_size)
np.random.shuffle(data)
model.load_dict(paddle.load("cvfo_model_rank.pdparams"))
model.eval()
batch_data = paddle.to_tensor(data[7]).astype('int64').reshape([1,3,-1])
star_seq = batch_data[:,:2,:9]
star_index=star_seq.shape[-1]-1
out_list=batch_data[:,2:,:9].reshape([-1]).numpy().tolist()
for i in tqdm(range(50)):
    if star_seq.shape[-1]%block_size != 0:
        star_seq = paddle.concat([star_seq, (len(voc)-1)*paddle.ones([1,2, 8-star_seq.shape[-1]%block_size]).astype("int64")], axis=-1)
    out = model(star_seq)
    if (star_index+1)%block_size==0:
        star_seq = paddle.concat([star_seq, (len(voc)-1)*paddle.ones([1,2, 8]).astype("int64")], axis=-1)
    out=voc.loc[voc["output_voc"].astype("int")==paddle.argmax(out, axis=-1)[:,star_index].numpy()[0]].values[0]
    star_seq[:,:,star_index+1]=paddle.to_tensor(out[1:3].astype("int")).astype('int64').reshape([1,2])
    star_index += 1
    out_list.append(out[-1])
print([voc.loc[voc["output_voc"].astype("int")==i,"voc"].values[0] for i in out_list])